from typing import TYPE_CHECKING

import torch.nn as nn
from rainforeLearn.gomoku.v2.neywork.blocks import NoisyLinear
if TYPE_CHECKING:
    from network import GomokuNetwork


class NetworkOutputHeads:
    """网络输出头构建器"""

    @staticmethod
    def build_multi_task_heads(network: 'GomokuNetwork', dropout_rate: float) -> None:
        """构建多任务输出头"""
        NetworkOutputHeads._build_action_head(network, dropout_rate)
        NetworkOutputHeads._build_value_head(network, dropout_rate)

    @staticmethod
    def build_single_task_head(network: 'GomokuNetwork', dropout_rate: float) -> None:
        """构建单任务输出头"""
        network.output_conv = nn.Conv2d(network.num_channels, 32, kernel_size=1, bias=False)
        network.output_bn = nn.BatchNorm2d(32)
        network.output_dropout = nn.Dropout(dropout_rate)

        linear_input_size = 32 * network.board_size ** 2

        if network.use_noisy:
            network.output_fc1 = NoisyLinear(linear_input_size, 1024)
            network.output_fc2 = NoisyLinear(1024, 512)
            network.output_fc3 = NoisyLinear(512, network.output_size)
        else:
            network.output_fc1 = nn.Linear(linear_input_size, 1024)
            network.output_fc2 = nn.Linear(1024, 512)
            network.output_fc3 = nn.Linear(512, network.output_size)

    @staticmethod
    def _build_action_head(network: 'GomokuNetwork', dropout_rate: float) -> None:
        """构建动作预测头"""
        network.action_conv = nn.Conv2d(network.num_channels, 16, kernel_size=1, bias=False)
        network.action_bn = nn.BatchNorm2d(16)
        network.action_dropout = nn.Dropout(dropout_rate)

        linear_input_size = 16 * network.board_size ** 2

        if network.use_noisy:
            network.action_fc1 = NoisyLinear(linear_input_size, 1024)
            network.action_fc2 = NoisyLinear(1024, 512)
            network.action_fc3 = NoisyLinear(512, network.action_size)
        else:
            network.action_fc1 = nn.Linear(linear_input_size, 1024)
            network.action_fc2 = nn.Linear(1024, 512)
            network.action_fc3 = nn.Linear(512, network.action_size)

    @staticmethod
    def _build_value_head(network: 'GomokuNetwork', dropout_rate: float) -> None:
        """构建价值评估头"""
        network.value_conv = nn.Conv2d(network.num_channels, 8, kernel_size=1, bias=False)
        network.value_bn = nn.BatchNorm2d(8)
        network.value_dropout = nn.Dropout(dropout_rate)

        linear_input_size = 8 * network.board_size ** 2

        if network.use_noisy:
            network.value_fc1 = NoisyLinear(linear_input_size, 512)
            network.value_fc2 = NoisyLinear(512, 128)
            network.value_fc3 = NoisyLinear(128, network.value_size)
        else:
            network.value_fc1 = nn.Linear(linear_input_size, 512)
            network.value_fc2 = nn.Linear(512, 128)
            network.value_fc3 = nn.Linear(128, network.value_size)
